import datetime
import backtrader as bt
import pandas as pd
import akshare as ak


# https://zhuanlan.zhihu.com/p/114782214
# https://zhuanlan.zhihu.com/p/380361696 order.executed.value的含义是该订单的成本
# 比如，你花100元买入1股股票，成本为100元，后来120卖掉它，成本还是100元，所以成本是不变的


def get_data(code):
    df = ak.stock_zh_a_daily(symbol=code, adjust='qfq')
    df.index = pd.to_datetime(df.date)
    df.sort_index(inplace=True)

    return df


# https://zhuanlan.zhihu.com/p/494061468 sizer的开发
class TestSizer(bt.Sizer):
    params = (('stake', 1),)

    def _getsizing(self, comminfo, cash, data, isbuy):
        if isbuy:
            return self.p.stake
        position = self.broker.getposition(data)
        if not position.size:
            return 0
        else:
            return position.size
        # print(f'XL={self.p.stake}')
        # return self.p.stake


class TestStrategy(bt.Strategy):
    params = (('maperiod', 15), ('printlog', False),)

    def log(self, txt, dt=None, doprint=False):
        if self.params.printlog or doprint:
            dt = dt or self.datas[0].datetime.date(0)
            print('%s, %s' % (dt.isoformat(), txt))

    def __init__(self):

        self.dataclose = self.datas[0].close
        self.datahigh = self.datas[0].high
        self.datalow = self.datas[0].low

        self.order = None
        self.buyprice = 0
        self.buycomm = 0
        self.newstake = 0
        self.buytime = 0
        # 参数计算，唐奇安通道上轨、唐奇安通道下轨、ATR
        self.DonchianHi = bt.indicators.Highest(self.datahigh(-1), period=20, subplot=False)
        self.DonchianLo = bt.indicators.Lowest(self.datalow(-1), period=10, subplot=False)
        self.TR = bt.indicators.Max((self.datahigh(0) - self.datalow(0)), abs(self.dataclose(-1) - self.datahigh(0)),
                                    abs(self.dataclose(-1) - self.datalow(0)))
        self.ATR = bt.indicators.SimpleMovingAverage(self.TR, period=self.p.maperiod, subplot=True)
        # 唐奇安通道上轨突破、唐奇安通道下轨突破
        self.CrossoverHi = bt.ind.CrossOver(self.dataclose(0), self.DonchianHi)
        self.CrossoverLo = bt.ind.CrossOver(self.dataclose(0), self.DonchianLo)

    def notify_order(self, order):
        if order.status in [order.Submitted, order.Accepted]:
            return

        if order.status in [order.Completed]:
            if order.isbuy():
                self.log(
                    'BUY EXECUTED, Price: %.2f, Cost: %.2f, Comm %.2f' %
                    (order.executed.price,
                     order.executed.value,
                     order.executed.comm), doprint=True)
                self.buyprice = order.executed.price
                self.buycomm = order.executed.comm
            else:
                self.log('SELL EXECUTED, Price: %.2f, Cost: %.2f, Comm %.2f' %
                         (order.executed.price,
                          order.executed.value,
                          order.executed.comm), doprint=True)
                self.bar_executed = len(self)
        elif order.status in [order.Canceled, order.Margin, order.Rejected]:
            self.log('Order Canceled/Margin/Rejected')
        self.order = None

    def notify_trade(self, trade):
        if not trade.isclosed:
            return
        self.log('OPERATION PROFIT, GROSS %.2f, NET %.2f' % (trade.pnl, trade.pnlcomm))

    def next(self):
        if self.order:
            return
            # 入场
        print(f"date={self.data.datetime.date(0)}->close={self.dataclose[0]}")
        if self.CrossoverHi > 0 and self.buytime == 0:
            self.newstake = self.broker.getvalue() * 0.01 / self.ATR
            print(
                f'date={self.data.datetime.date(0)}->atr={self.ATR[0]}->value={self.broker.getvalue() * 0.01}->stake={self.broker.getvalue() * 0.01 / self.ATR}')
            self.newstake = int(self.newstake / 100) * 100
            self.sizer.p.stake = self.newstake
            # self.sizer.p.stake = 100
            self.buytime = 1
            self.order = self.buy()
            print(f"1->size={self.sizer.p.stake}")
            # 加仓
        elif self.datas[0].close > self.buyprice + 0.5 * self.ATR[0] and 0 < self.buytime < 3:
            self.newstake = self.broker.getvalue() * 0.01 / self.ATR
            self.newstake = int(self.newstake / 100) * 100
            self.sizer.p.stake = self.newstake
            # self.sizer.p.stake = 100
            self.order = self.buy()
            self.buytime = self.buytime + 1
            # 出场
        elif self.CrossoverLo < 0 and self.buytime > 0:
            self.order = self.sell()  # TestSizer已经重载卖出的size，当卖出时，清仓卖出
            self.buytime = 0
            # 止损
        elif self.datas[0].close < (self.buyprice - 2 * self.ATR[0]) and self.buytime > 0:
            self.order = self.sell()
            self.buytime = 0

    def stop(self):
        self.log('(MA Period %2d) Ending Value %.2f' % (self.params.maperiod, self.broker.getvalue()), doprint=True)


if __name__ == '__main__':
    code = 'sz300274'
    st_date = datetime.datetime(2023, 1, 1)
    ed_date = datetime.datetime(2023, 6, 21)
    data = get_data(code)
    datafeed = bt.feeds.PandasData(dataname=data, fromdate=st_date, todate=ed_date)
    cerebro = bt.Cerebro()
    cerebro.addstrategy(TestStrategy)
    cerebro.adddata(datafeed, name=code)
    cerebro.broker.setcash(100000.0)
    cerebro.broker.setcommission(commission=0.001)
    # 设置买入策略
    cerebro.addsizer(TestSizer)
    print('Starting Portfolio Value: %.2f' % cerebro.broker.getvalue())
    # 启动回测
    cerebro.run()
    print('Final Portfolio Value: %.2f' % cerebro.broker.getvalue())
    # 曲线绘图输出
    cerebro.plot()
